import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler

# 添加一个全局标准化器，用于保存训练数据的标准化参数
_scaler = None
_last_values = None  # 保存训练数据的最后几个值用于预测

def prepare_features(df, is_predict=False):
    """组合多种特征"""
    global _scaler, _last_values
    
    dates = pd.to_datetime(df.iloc[:, 0])
    year = dates.dt.year
    month = dates.dt.month
    
    # 季节性
    month_sin = np.sin(2 * np.pi * month/12)
    month_cos = np.cos(2 * np.pi * month/12)
    quarter = dates.dt.quarter
    # 是否节假日
    is_holiday = np.where(month.isin([1, 2, 7, 8, 10]), 1, 0)
    # 温度特征
    temp_proxy = -((month - 7) ** 2) + 49
    
    # 滑动窗口特征
    window_size = 3
    values = df.iloc[:, 1].values 
    if not is_predict:
        # 训练模式：计算滑动窗口特征
        rolling_series = pd.Series(values)
        rolling_mean = rolling_series.rolling(window=window_size, min_periods=1).mean()
        rolling_mean = rolling_mean.bfill()
        rolling_std = rolling_series.rolling(window=window_size, min_periods=1).std().fillna(0)
        # 保存最后几个值用于预测
        _last_values = values[-window_size:]
    else:
        # 预测模式：使用保存的历史数据
        if _last_values is None:
            # 如果没有历史数据，用默认值
            rolling_mean = np.array([np.mean(values)] * len(values))
            rolling_std = np.array([0] * len(values))
        else:
            # 历史数据的均值和标准差
            rolling_mean = np.array([np.mean(_last_values)] * len(values))
            rolling_std = np.array([np.std(_last_values)] * len(values))
    
    # 趋势特征：使用固定的起始日期
    base_date = pd.Timestamp('2020-01-01')
    months_since_start = ((dates.dt.year - base_date.year) * 12 + 
                         (dates.dt.month - base_date.month))
    
    # 组合特征
    X = np.column_stack([
        year,
        month,
        month_sin,
        month_cos,
        quarter,
        is_holiday,
        temp_proxy,
        rolling_mean,
        rolling_std,
        months_since_start
    ])
    
    # 标准化处理
    if not is_predict:
        # 训练模式：拟合并转换
        _scaler = StandardScaler()
        X = _scaler.fit_transform(X)
        y = df.iloc[:, 1].values
        return X, y
    else:
        # 预测模式：使用已有的标准化参数
        if _scaler is None:
            raise ValueError("模型未训练，请先进行训练")
        X = _scaler.transform(X)
        return X